import numpy as np

class TreeNode:
    def __init__(self,label):
        #分支节点存储分支features，叶节点存储结果标记
        self.label = label
        self.childs = {}


class Decision_Tree:
    def __init__(self):
        pass

    def build_tree(self,data_set,label_name,feature_names):
        labels = data_set[label_name]
        #完全纯度
        if np.unique(labels).shape[0]==1:
            return TreeNode(labels.values[0])
        #只剩最终判定特征(只剩一列特征)
        if data_set.shape[1]==1:
            result = self.get_main_label(data_set[label_name])
            return TreeNode(result)
        
        #print("当前数据集:\n",data_set)
        #print("当前特征:\n",feature_names)
        best_label = self.get_best_label(data_set,label_name,feature_names)
        #print("最好标签:",best_label)

        t = TreeNode(best_label)
        unique_values = np.unique(data_set[best_label])
        for unique_value in unique_values:
            new_data_set = data_set[data_set[best_label]==unique_value]
            new_data_set = new_data_set.drop(columns=[best_label])
            new_feature_names = [x for x in feature_names if x != best_label]
            child = self.build_tree(new_data_set,label_name,new_feature_names)
            t.childs[unique_value] = child
        return t

    def get_main_label(self,labels):
        print(labels.values)
        max_cnt_label = labels.values[0]
        max_cnt = 1
        for unique_label in np.unique(labels):
            n = labels[labels==unique_label].shape[0]
            if n>max_cnt:
                max_cnt_label = unique_label
                max_cnt = n
        return max_cnt_label
    
    def get_best_label(self,data_set,label_name,feature_names):
        best_label = feature_names[0]
        best_InfoGain = 0
        baseEntropy = self.compute_entroy(data_set,label_name)
        print("baseEntropy:",baseEntropy)
        for feature_name in feature_names:
            branchEntropy = 0
            sample_cnt = data_set.shape[0]
            unique_values = np.unique(data_set[feature_name])
            for unique_value in unique_values:
                part_sample = data_set[data_set[feature_name]==unique_value]
                part_sample_cnt = part_sample.shape[0]
                branchEntropy += (part_sample_cnt/sample_cnt)*self.compute_entroy(part_sample,label_name)
            infoGain = baseEntropy - branchEntropy
            if infoGain>best_InfoGain:
                best_InfoGain = infoGain
                best_label = feature_name
        return best_label
            
    def compute_entroy(self,data_set,label_name):
        unique_values = np.unique(data_set[label_name])
        total_cnt = data_set.shape[0]
        entropy = 0
        for unique_value in unique_values:
            cnt = data_set[data_set[label_name]==unique_value].shape[0]
            prob = cnt/total_cnt
            entropy += -prob*np.log2(prob)
        return entropy
    
    def classify(self,tree,x,feature_names):
        if len(tree.childs)==0:
            return tree.label
        feature_name = tree.label
        feature_value = x[feature_names.index(feature_name)]
        if feature_value in tree.childs.keys(): 
            return self.classify(tree.childs[feature_value],x,feature_names)
        return tree.label


